from fastapi import FastAPI, BackgroundTasks, WebSocket, WebSocketDisconnect, APIRouter
from fastapi.responses import JSONResponse
import time
import asyncio


task_router = APIRouter()


# 用于存储任务进度和日志的全局变量
task_progress = {}
task_logs = {}

def background_task(task_id: str):
    for i in range(1, 11):
        time.sleep(1)  # 模拟长时间运行的任务
        task_progress[task_id] = i * 10  # 更新任务进度
        task_logs[task_id].append(f"Task {task_id} progress: {i * 10}%")  # 添加日志
    task_progress[task_id] = 100  # 任务完成
    task_logs[task_id].append(f"Task {task_id} completed")

@task_router.post("/start-task")
def start_task(background_tasks: BackgroundTasks):
    task_id = str(time.time())  # 使用时间戳作为任务ID
    task_progress[task_id] = 0  # 初始化任务进度
    task_logs[task_id] = []  # 初始化任务日志
    background_tasks.add_task(background_task, task_id)
    return {"task_id": task_id}

@task_router.websocket("/ws/{task_id}")
async def websocket_endpoint(websocket: WebSocket, task_id: str):
    await websocket.accept()
    try:
        while True:
            progress = task_progress.get(task_id, 0)
            logs = task_logs.get(task_id, [])
            await websocket.send_json({"type": "progress", "progress": progress})
            for log in logs:
                await websocket.send_json({"type": "log", "log": log})
            await asyncio.sleep(1)
    except WebSocketDisconnect:
        print(f"WebSocket connection closed for task {task_id}")


